from osgeo import gdal
from osgeo import gdalconst
import zipfile
import shutil
import sys
import os
import re


def unzip_l1c(zip_path, bands, save_dir):
    band_names = {
        "1":"T.*_B01\.jp2",
        "2":"T.*_B02\.jp2",
        "3":"T.*_B03\.jp2",
        "4":"T.*_B04\.jp2",
        "5":"T.*_B05\.jp2",
        "6":"T.*_B06\.jp2",
        "7":"T.*_B07\.jp2",
        "8A":"T.*_B8A\.jp2",
        "8":"T.*_B08\.jp2",
        "9":"T.*_B09\.jp2",
        "10":"T.*_B10\.jp2",
        "11":"T.*_B11\.jp2",
        "12":"T.*_B12\.jp2"}
    pathres = list(map(lambda x:".*/IMG_DATA/"+band_names[f"{x}"],bands))
    paths = []
    zipFile = zipfile.ZipFile(zip_path, "r")
    for pathre in pathres:
        for info in zipFile.infolist():
            if re.match(pathre,info.filename):
                path = zipFile.extract(info.filename, save_dir)
                paths.append(path)
    
    zipFile.close()
    return paths


def core(jp2_paths,tif_path):
    vrt_path = tif_path.replace('.tif','.vrt')
    options = gdal.BuildVRTOptions(
        resolution="highest",
        resampleAlg=gdalconst.GRA_NearestNeighbour,
        separate=True)
    gdal.BuildVRT(vrt_path,jp2_paths,options=options)
    creationOptions = [
        "BIGTIFF=YES",
        "TILED=YES",
        "BLOCKXSIZE=1024",
        # "TFW=YES",
        "NUM_THREADS=ALL_CPUS"]
    options = gdal.TranslateOptions(
        format="GTiff",
        xRes=10,
        yRes=10,
        resampleAlg=gdalconst.GRA_NearestNeighbour,
        creationOptions=creationOptions)
    gdal.Translate(tif_path,vrt_path,options=options)
    os.remove(vrt_path)
    return


def get_l1c_raster(zip_path, tif_path, bands):
    # zip_name = os.path.basename(zip_path)
    # year_month = zip_name.split("_")[2][:6]
    # date_dir = os.path.join(save_dir,year_month)
    date_dir = os.path.dirname(tif_path)
    if not os.path.exists(date_dir):os.makedirs(date_dir)
    # raster_name = zip_name.replace('.zip','.tif')
    # raster_path = os.path.join(date_dir, raster_name)
    if os.path.exists(tif_path): return tif_path
    jp2_paths = unzip_l1c(zip_path,bands,date_dir)
    core(jp2_paths,tif_path)
    name = os.path.basename(zip_path).replace(".zip",".SAFE")
    unzip_dir = os.path.join(date_dir,name)
    shutil.rmtree(unzip_dir)
    return


def main():
    return


if __name__ == "__main__":
    zip_path = r"/data/RS_Data/Image_Data/Sentinel2/1C/202403/S2B_MSIL1C_20240314T025519_N0510_R032_T50SNJ_20240314T044704.zip"
    save_dir = r"/data/Yang/dsafa"
    bands = ['2','3','4','5','6','7',"8A",'8','11','12']
    get_l1c_raster(zip_path, save_dir, bands)
    # get_l1c_raster(sys.argv[1], sys.argv[2], bands)

